#include <ros/ros.h>
#include <sensor_msgs/NavSatFix.h>
#include <geometry_msgs/Vector3Stamped.h>
#include <geometry_msgs/QuaternionStamped.h>
#include <nav_msgs/Odometry.h>
#include <nav_msgs/Path.h>
#include "transform_utils/geometry_transform.h"
#include "transform_utils/type_transform.h"
#include "transform_utils/display.h"
#include "std_msgs/Bool.h"
#include <rosbag/bag.h>

#include <Eigen/Core>
#include <Eigen/Geometry>
#include <Eigen/Dense>

using Eigen::MatrixXd;
using Eigen::Quaterniond;
using Eigen::Vector3d;
using std::string;
using std::vector;
using namespace Planner;

bool laser_or_gps, init_gps = false, init_transform = false, record_flag = true;
double init_latitude, init_longitude, init_height;
double last_lautitude, last_longitude;
Vector3d gps_pos, laser_pos;
Quaterniond dji_orien, laser_orien, flu2enu;
double C_EARTH = 6378137.0;

void gpsPositionCallback(const sensor_msgs::NavSatFix &msg)
{
    // 计算结果为ENU坐标系
    if (!init_gps)
    {
        init_latitude = msg.latitude;
        init_longitude = msg.longitude;
        init_height = msg.altitude;
        last_lautitude = msg.latitude;
        last_longitude = msg.longitude;
        init_gps = true;
        ROS_WARN("Initial height: %f", init_height);
    }
    else
    {
        double delta_x = toRad(msg.latitude - last_lautitude) * C_EARTH;
        double delta_y = toRad(msg.longitude - last_longitude) * C_EARTH * cos(toRad(msg.latitude));
        // 异常处理，有的时候gps信号会丢失
        if (abs(delta_x) > 10 || abs(delta_y) > 10)
            return;
        gps_pos[1] = toRad(msg.latitude - init_latitude) * C_EARTH;
        gps_pos[0] = toRad(msg.longitude - init_longitude) * C_EARTH * cos(toRad(msg.latitude));
        gps_pos[2] = msg.altitude - init_height;
        last_longitude = msg.longitude;
        last_lautitude = msg.latitude;
        ROS_INFO("GPS INFO: latitude: %f, longitude: %f", msg.latitude, msg.longitude);
    }
}

void djiOrienCallback(const geometry_msgs::QuaternionStamped &msg)
{
    // 姿态角为FLU坐标系相对于ENU的初始姿态
    geometry2EigenQuaternion(msg.quaternion, dji_orien);

    if (!init_transform)
    {
        flu2enu = dji_orien;
        init_transform = true;
    }
}

void laserOdomCallback(const nav_msgs::Odometry &msg)
{
    geometry2EigenVector(msg.pose.pose.position, laser_pos);
    geometry2EigenQuaternion(msg.pose.pose.orientation, laser_orien);
}

void stopRecordCallback(const std_msgs::Bool &msg)
{
    record_flag = false;
}

void savePathToBag(const vector<Vector3d> &path, string file_name);

int main(int argc, char *argv[])
{
    ros::init(argc, argv, "record_manual_flight");
    ros::NodeHandle nh;

    ros::Subscriber gps_position_sub;
    ros::Subscriber dji_orien_sub;
    ros::Subscriber laser_odom_sub;
    ros::Subscriber stop_record_sub;
    ros::Publisher record_traj_pub;
    ros::Publisher record_odom_pub;

    string odom_topic, bag_file;
    nh.param("record_manual_flight/laser_odom_topic", odom_topic, string(""));
    nh.param("record_manual_flight/save_bag_name", bag_file, string(""));
    nh.param("record_manual_flight/laser_or_gps", laser_or_gps, false);

    gps_position_sub = nh.subscribe("/dji_osdk_ros/gps_position", 1, gpsPositionCallback);
    dji_orien_sub = nh.subscribe("/dji_osdk_ros/attitude", 1, djiOrienCallback);
    laser_odom_sub = nh.subscribe(odom_topic, 1, laserOdomCallback);
    stop_record_sub = nh.subscribe("/finish_trigger", 1, stopRecordCallback);
    record_traj_pub = nh.advertise<nav_msgs::Path>("record_traj_vis", 1);
    record_odom_pub = nh.advertise<nav_msgs::Odometry>("record_odom_vis", 1);

    vector<Vector3d> path;
    ros::Rate rate(50);
    while (ros::ok())
    {
        ros::spinOnce();

        if (!init_gps) {
            rate.sleep();
            continue;
        }

        Vector3d curr_pos;
        if (laser_or_gps)
        {
            // laser data
            curr_pos = laser_pos;
        }
        else
        {
            // gps and dji data
            curr_pos = flu2enu.inverse() * gps_pos;
        }

        if (path.empty())
        {
            if (curr_pos[2] > 1)
            {
                path.push_back(curr_pos);
                ROS_INFO("Start record");
            }
        }
        else if ((curr_pos - path.back()).norm() > 1.5)
        {
            path.push_back(curr_pos);
        }

        nav_msgs::Odometry odom_msg;
        odom_msg.header.frame_id = "world";
        odom_msg.pose.pose.position.x = curr_pos[0];
        odom_msg.pose.pose.position.y = curr_pos[1];
        odom_msg.pose.pose.position.z = curr_pos[2];
        record_odom_pub.publish(odom_msg);

        if (!record_flag)
        {
            savePathToBag(path, bag_file);
            break;
        }

        rate.sleep();
    }

    for (size_t i = 0; i < 10; i++)
    {
        displayTrajectory(path, record_traj_pub);
        rate.sleep();
    }

    return 0;
}

void savePathToBag(const vector<Vector3d> &path, string file_name)
{
    nav_msgs::Path msgs;
    msgs.header.frame_id = "world";
    geometry_msgs::PoseStamped msg;
    msg.header.frame_id = "world";

    for (auto p : path)
    {
        msg.pose.position.x = p[0];
        msg.pose.position.y = p[1];
        msg.pose.position.z = p[2];
        msgs.poses.push_back(msg);
    }
    rosbag::Bag bag;
    bag.open(file_name, rosbag::bagmode::Write);
    bag.write("record_trag", ros::Time::now(), msgs);
    ROS_INFO("Write path to rosbag");
}
